# pip install  datasets==2.18

import os, json
from transformers import BertTokenizer
from modelscope.models import Model
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.preprocessors import TableQuestionAnsweringPreprocessor
from modelscope.preprocessors.nlp.space_T_cn.fields.database import Database
from modelscope.utils.constant import ModelFile, Tasks
import logging

logger = logging.getLogger(__name__)


class Text2Sql:
    """Class to connect to the Rasa X API"""

    def __init__(self ):
        print('init')

    def init(self ):
        self.model = Model.from_pretrained('/home/softrobot/webapps/bot/workflow/actions/api/damo/nlp_convai_text2sql_pretrain_cn')
        print('model.model_dir='+self.model.model_dir)

        self.tokenizer = BertTokenizer(os.path.join(self.model.model_dir, ModelFile.VOCAB_FILE))
        self.db = Database(
            tokenizer=self.tokenizer,
            table_file_path=os.path.join(self.model.model_dir, 'table.json'),
            syn_dict_file_path=os.path.join(self.model.model_dir, 'synonym.txt'),
            is_use_sqlite=True)
        self.preprocessor = TableQuestionAnsweringPreprocessor(model_dir=self.model.model_dir, db=self.db)
        self.pipelines = [
            pipeline(
                Tasks.table_question_answering,
                model=self.model,
                preprocessor=self.preprocessor,
                db=self.db)
        ]


    def toSql(self ,test_case):
        for pipeline in self.pipelines:
            historical_queries = None
            for question, table_id in test_case['sql']:
                output_dict = pipeline({
                    'question': question,
                    'table_id': table_id,
                    'history_sql': historical_queries
                })[OutputKeys.OUTPUT]
                print('question', question)
                print('sql text:', output_dict[OutputKeys.SQL_STRING])
                # logger.info('sql query:', output_dict[OutputKeys.SQL_QUERY])
                ret = output_dict[OutputKeys.SQL_QUERY]
                ret = ret.replace('==', '=')
                ret = ret.replace('!=', '<>')
                if '.null' in ret:
                    return '没有权限或找不到对应表结构'
                return ret


sql = Text2Sql()
test_case = {
    'sql':
        [
            ['[sql]长江流域的小型水库的库容总量是多少？', ''],
            ['[sql]列出油耗大于8但是功率低于200的名称和价格？', ''],
            ['[sql]计算机或者成绩优秀的同学有哪些？学号是多少？', ''],
            ['[sql]上个月收益超过3的有几个基金？', ''],
            ['[sql]油耗低于5的suv有哪些？', '']
        ]
}
sql.init()
ret = sql.toSql(test_case)
logger.info('###### ret =')